package ar.com.hjg.pngj;

import java.util.zip.DataFormatException;
import java.util.zip.Inflater;

/**
 * A set of IDAT-like chunks which, concatenated, form a zlib stream.
 * <p>
 * The inflated stream is intented to be read as a sequence of "rows", of which the caller knows the lengths (not necessary equal) and number.
 * <p>
 * Eg: For IDAT non-interlaced images, a row has bytesPerRow + 1 filter byte<br>
 * For interlaced images, the lengths are variable.
 * <p>
 * This class can work in sync (polled) mode or async (callback) mode. But for callback mode the method processRowCallback() must be overriden
 * <p>
 * See {@link IdatSet}, which is mostly used and has a slightly simpler use.<br>
 * See <code>DeflatedChunkSetTest</code> for example of use.
 */
public class DeflatedChunksSet {

  protected byte[] row; // a "row" here means a raw (uncopressed filtered) part of the IDAT stream,
                        // normally a image row (or subimage row for interlaced) plus a filter byte
  private int rowfilled; // effective/valid length of row
  private int rowlen; // what amount of bytes is to be interpreted as a complete "row". can change
                      // (for interlaced)
  private int rown; // only coincide with image row if non-interlaced - incremented by
                    // setNextRowSize()

  /*
   * States WAITING_FOR_INPUT ROW_READY WORK_DONE TERMINATED
   * 
   * processBytes() is externally called, prohibited in READY (in DONE it's ignored)
   * 
   * WARNING: inflater.finished() != DONE (not enough, not neccesary) DONE means that we have already uncompressed all the data of interest.
   * 
   * In non-callback mode, prepareForNextRow() is also externally called, in
   * 
   * Flow: - processBytes() calls inflateData() - inflateData() : if buffer is filled goes to READY else if ! inf.finished goes to WAITING else if any data goes to READY
   * (incomplete data to be read) else goes to DONE - in Callback mode, after going to READY, n=processCallback() is called and then prepareForNextRow(n) is called. - in Polled
   * mode, prepareForNextRow(n) must be called from outside (after checking state=READY) - prepareForNextRow(n) goes to DONE if n==0 calls inflateData() again - end() goes to DONE
   */
  private enum State {
    WAITING_FOR_INPUT, // waiting for more input
    ROW_READY, // ready for consumption (might be less than fully filled), ephemeral for CALLBACK
               // mode
    WORK_DONE, // all data of interest has been read, but we might accept still more trailing chunks
               // (we'll ignore them)
    TERMINATED; // we are done, and also won't accept more IDAT chunks

    public boolean isDone() {
      return this == WORK_DONE || this == TERMINATED;
    } // the caller has already uncompressed all the data of interest or EOF

    public boolean isTerminated() {
      return this == TERMINATED;
    } // we dont accept more chunks
  }

  State state = State.WAITING_FOR_INPUT; // never null

  private Inflater inf;
  private final boolean infOwn; // true if we own the inflater (we created it)

  private DeflatedChunkReader curChunk;

  private boolean callbackMode = true;
  private long nBytesIn = 0; // count the total compressed bytes that have been fed
  private long nBytesOut = 0; // count the total uncompressed bytes
  int chunkNum = -1; // incremented at each new chunk start
  int firstChunqSeqNum = -1; // expected seq num for first chunk. used only for fDAT (APNG)

  /**
   * All IDAT-like chunks that form a same DeflatedChunksSet should have the same id
   */
  public final String chunkid;

  /**
   * @param initialRowLen Length in bytes of first "row" (see description)
   * @param maxRowLen Max length in bytes of "rows"
   * @param inflater Can be null. If not null, must be already reset (and it must be closed/released by caller!)
   */
  public DeflatedChunksSet(String chunkid, int initialRowLen, int maxRowLen, Inflater inflater,
      byte[] buffer) {
    this.chunkid = chunkid;
    this.rowlen = initialRowLen;
    if (initialRowLen < 1 || maxRowLen < initialRowLen)
      throw new PngjException("bad inital row len " + initialRowLen);
    if (inflater != null) {
      this.inf = inflater;
      infOwn = false;
    } else {
      this.inf = new Inflater();
      infOwn = true; // inflater is own, we will release on close()
    }
    this.row = buffer != null && buffer.length >= initialRowLen ? buffer : new byte[maxRowLen];
    rown = -1;
    this.state = State.WAITING_FOR_INPUT;
    try {
      prepareForNextRow(initialRowLen);
    } catch (RuntimeException e) {
      close();
      throw e;
    }
  }

  public DeflatedChunksSet(String chunkid, int initialRowLen, int maxRowLen) {
    this(chunkid, initialRowLen, maxRowLen, null, null);
  }

  protected void appendNewChunk(DeflatedChunkReader cr) {
    // all chunks must have same id
    if (!this.chunkid.equals(cr.getChunkRaw().id))
      throw new PngjInputException("Bad chunk inside IdatSet, id:" + cr.getChunkRaw().id
          + ", expected:" + this.chunkid);
    this.curChunk = cr;
    chunkNum++;
    if (firstChunqSeqNum >= 0)
      cr.setSeqNumExpected(chunkNum + firstChunqSeqNum);
  }

  /**
   * Feeds the inflater with the compressed bytes
   * 
   * In poll mode, the caller should not call repeatedly this, without consuming first, checking isDataReadyForConsumer()
   * 
   * @param buf
   * @param off
   * @param len
   */
  protected void processBytes(byte[] buf, int off, int len) {
    nBytesIn += len;
    // PngHelperInternal.LOGGER.info("processing compressed bytes in chunkreader : " + len);
    if (len < 1 || state.isDone())
      return;
    if (state == State.ROW_READY)
      throw new PngjInputException("this should only be called if waitingForMoreInput");
    if (inf.needsDictionary() || !inf.needsInput())
      throw new RuntimeException("should not happen");
    inf.setInput(buf, off, len);
    // PngHelperInternal.debug("entering processs bytes, state=" + state +
    // " callback="+callbackMode);
    if (isCallbackMode()) {
      while (inflateData()) {
        int nextRowLen = processRowCallback();
        prepareForNextRow(nextRowLen);
        if (isDone())
          processDoneCallback();
      }
    } else
      inflateData();
  }

  /*
   * This never inflates more than one row This returns true if this has resulted in a row being ready and preprocessed with preProcessRow (in callback mode, we should call
   * immediately processRowCallback() and prepareForNextRow(nextRowLen)
   */
  private boolean inflateData() {
    try {
      // PngHelperInternal.debug("entering inflateData bytes, state=" + state +
      // " callback="+callbackMode);
      if (state == State.ROW_READY)
        throw new PngjException("invalid state");// assert
      if (state.isDone())
        return false;
      int ninflated = 0;
      if (row == null || row.length < rowlen)
        row = new byte[rowlen]; // should not happen
      if (rowfilled < rowlen && !inf.finished()) {
        try {
          ninflated = inf.inflate(row, rowfilled, rowlen - rowfilled);
        } catch (DataFormatException e) {
          throw new PngjInputException("error decompressing zlib stream ", e);
        }
        rowfilled += ninflated;
        nBytesOut += ninflated;
      }
      State nextstate = null;
      if (rowfilled == rowlen)
        nextstate = State.ROW_READY; // complete row, process it
      else if (!inf.finished())
        nextstate = State.WAITING_FOR_INPUT;
      else if (rowfilled > 0)
        nextstate = State.ROW_READY; // complete row, process it
      else {
        nextstate = State.WORK_DONE; // eof, no more data
      }
      state = nextstate;
      if (state == State.ROW_READY) {
        preProcessRow();
        return true;
      }
    } catch (RuntimeException e) {
      close();
      throw e;
    }
    return false;
  }

  /**
   * Called automatically in all modes when a full row has been inflated.
   */
  protected void preProcessRow() {

  }

  /**
   * Callback, must be implemented in callbackMode
   * <p>
   * This should use {@link #getRowFilled()} and {@link #getInflatedRow()} to access the row.
   * <p>
   * Must return byes of next row, for next callback.
   */
  protected int processRowCallback() {
    throw new PngjInputException("not implemented");
  }

  /**
   * Callback, to be implemented in callbackMode
   * <p>
   * This will be called once to notify state done
   */
  protected void processDoneCallback() {}

  /**
   * Inflated buffer.
   * 
   * The effective length is given by {@link #getRowFilled()}
   */
  public byte[] getInflatedRow() {
    return row;
  }

  /**
   * Should be called after the previous row was processed
   * <p>
   * Pass 0 or negative to signal that we are done (not expecting more bytes)
   * <p>
   * This resets {@link #rowfilled}
   * <p>
   * The
   */
  public void prepareForNextRow(int len) {
    rowfilled = 0;
    rown++;
    if (len < 1) {
      rowlen = 0;
      done();
    } else if (inf.finished()) {
      rowlen = 0;
      done();
    } else {
      state = State.WAITING_FOR_INPUT;
      rowlen = len;
      if (!callbackMode)
        inflateData();
    }
  }

  /**
   * In this state, the object is waiting for more input to deflate.
   * <p>
   * Only in this state it's legal to feed this
   */
  public boolean isWaitingForMoreInput() {
    return state == State.WAITING_FOR_INPUT;
  }

  /**
   * In this state, the object is waiting the caller to retrieve inflated data
   * <p>
   * Effective length: see {@link #getRowFilled()}
   */
  public boolean isRowReady() {
    return state == State.ROW_READY;
  }

  /**
   * In this state, all relevant data has been uncompressed and retrieved (exceptionally, the reading has ended prematurely).
   * <p>
   * We can still feed this object, but the bytes will be swallowed/ignored.
   */
  public boolean isDone() {
    return state.isDone();
  }

  public boolean isTerminated() {
    return state.isTerminated();
  }

  /**
   * This will be called by the owner to report us the next chunk to come. We can make our own internal changes and checks. This returns true if we acknowledge the next chunk as
   * part of this set
   */
  public boolean ackNextChunkId(String id) {
    if (state.isTerminated())
      return false;
    else if (id.equals(chunkid)) {
      return true;
    } else {
      if (!allowOtherChunksInBetween(id)) {
        if (state.isDone()) {
          if (!isTerminated())
            terminate();
          return false;
        } else {
          throw new PngjInputException("Unexpected chunk " + id + " while " + chunkid
              + " set is not done");
        }
      } else
        return true;
    }
  }

  protected void terminate() {
    close();
  }

  /**
   * This should be called when discarding this object, or for aborting. Secure, idempotent Don't use this just to notify this object that it has no more work to do, see
   * {@link #done()}
   * */
  public void close() {
    try {
      if (!state.isTerminated()) {
        state = State.TERMINATED;
      }
      if (infOwn && inf != null) {
        inf.end();// we end the Inflater only if we created it
        inf = null;
      }
    } catch (Exception e) {
    }
  }

  /**
   * Forces the DONE state, this object won't uncompress more data. It's still not terminated, it will accept more IDAT chunks, but will ignore them.
   */
  public void done() {
    if (!isDone())
      state = State.WORK_DONE;
  }

  /**
   * Target size of the current row, including filter byte. <br>
   * should coincide (or be less than) with row.length
   */
  public int getRowLen() {
    return rowlen;
  }

  /** This the amount of valid bytes in the buffer */
  public int getRowFilled() {
    return rowfilled;
  }

  /**
   * Get current (last) row number.
   * <p>
   * This corresponds to the raw numeration of rows as seen by the deflater. Not the same as the real image row, if interlaced.
   * 
   */
  public int getRown() {
    return rown;
  }

  /**
   * Some IDAT-like set can allow other chunks in between (APGN?).
   * <p>
   * Normally false.
   * 
   * @param id Id of the other chunk that appeared in middel of this set.
   * @return true if allowed
   */
  public boolean allowOtherChunksInBetween(String id) {
    return false;
  }

  /**
   * Callback mode = async processing
   */
  public boolean isCallbackMode() {
    return callbackMode;
  }

  public void setCallbackMode(boolean callbackMode) {
    this.callbackMode = callbackMode;
  }

  /** total number of bytes that have been fed to this object */
  public long getBytesIn() {
    return nBytesIn;
  }

  /** total number of bytes that have been uncompressed */
  public long getBytesOut() {
    return nBytesOut;
  }

  @Override
  public String toString() {
    StringBuilder sb =
        new StringBuilder("idatSet : " + curChunk.getChunkRaw().id + " state=" + state + " rows="
            + rown + " bytes=" + nBytesIn + "/" + nBytesOut);
    return sb.toString();
  }

}
